import os

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # 去警告
# os.environ["CUDA_VISIBLE_DEVICES"] = "-1" # 使用CPU

import tensorflow as tf
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import matplotlib.pyplot as plt
import tensorflow.keras as keras


# %matplotlib inline  # 将图像内嵌进 notebook 中，PyCharm 不支持

# 过拟合：acc > val_acc，说明受高此项影响，调：
#     数据变换处理器，train_datagen = ImageDataGenerator
#     keras.layers.Dropout(0.5),
# x多分类调：
#     class_mode='categorical'，   - 两类：binary
#     keras.layers.Dense(x, activation='softmax'),   - 两类：activation='sigmoid'
#     损失函数：loss='categorical_crossentropy', （sparse_categorical_crossentropy：真实的标签值 y_true 可以直接传入 int 类型）
#                  - 两类：loss='binary_crossentropy'


if __name__ == "__main__":
    # 1. 数据准备
    # 1.1 数据所在文件夹
    base_dir = '../resources/p02_deep_learning_tensorFlow/dog_and_cat/'
    train_dir = os.path.join(base_dir, 'train/')
    validation_dir = os.path.join(base_dir, 'validation/')

    # # 训练集
    # train_cats_dir = os.path.join(train_dir, 'cats')
    # train_dogs_dir = os.path.join(train_dir, 'dogs')
    #
    # # 验证集
    # validation_cats_dir = os.path.join(validation_dir, 'cats')
    # validation_dogs_dir = os.path.join(validation_dir, 'dogs')

    # 2 数据预处理
    # 数据变换处理器 - 从数据角度解决过拟合问题
    train_datagen = ImageDataGenerator(
        rescale=1. / 255,             # 归一化
        rotation_range=40,            # 旋转角度
        width_shift_range=0.2,        # 横平移变换
        height_shift_range=0.2,       # 平纵移变换
        shear_range=0.2,              # 剪切范围
        zoom_range=0.2,               # 缩放
        horizontal_flip=True,         # 翻转
        fill_mode="nearest",
    )
    # 测试集不能数据变换
    valid_datagen = ImageDataGenerator(rescale=1. / 255)


    # 图像处理数据， GPU 可以 224*224*3
    width = 64
    height = 64
    channel = 3
    batch_size = 20  # 训练集每次迭代20个数据
    valid_batch_size = 20  # 测试集每次迭代20个数据
    # 2.1 从文件夹读取数据，并且每次读 20 个
    train_generator = train_datagen.flow_from_directory(
        train_dir,  # 文件夹路径
        target_size=(width, height),  # 指定resize成的大小，需要与模型的输入大小一致
        batch_size=batch_size,  # 每次迭代20个数据
        seed=7,  # 随机种子
        # 如果one-hot就是categorical，二分类用binary就可以
        class_mode='binary',
    )  # 数据格式
    # 查看数据：train_generator.next()

    validation_generator = valid_datagen.flow_from_directory(
        validation_dir,
        target_size=(width, height),
        batch_size=valid_batch_size,
        seed=7,
        class_mode='binary',
    )


    # 3. 特征工程
    # 估计器：基于 Keras 构建网络模型
    # 3.1 前向传播：按顺序构造网络模型
    model = keras.models.Sequential([
        # 如果训练慢，可以把数据设置的更小一些
        keras.layers.Conv2D(filters=32, kernel_size=3, activation='relu', padding='same', input_shape=[width, height, channel]),  # 卷积：生成64个特征图，卷积核 3*3，激活函数(概念化)，填充，输入 64*64*3
        keras.layers.MaxPooling2D(pool_size=2),                                             # 池化 - 最大池化：变成原来 1/2
        keras.layers.Conv2D(filters=32, kernel_size=3, activation='relu', padding='same'),  # 卷积：生成32个特征图，卷积核 3*3，激活函数(概念化),
        keras.layers.MaxPooling2D(pool_size=2),                                             # 池化

        keras.layers.Conv2D(filters=64, kernel_size=3, activation='relu', padding='same'),  # 卷积
        keras.layers.MaxPooling2D(pool_size=2),                                             # 池化
        keras.layers.Conv2D(filters=64, kernel_size=3, activation='relu', padding='same'),  # 卷积
        keras.layers.MaxPooling2D(pool_size=2),                                             # 池化

        keras.layers.Conv2D(filters=128, kernel_size=3, activation='relu', padding='same'), # 卷积
        keras.layers.MaxPooling2D(pool_size=2),                                             # 池化
        keras.layers.Conv2D(filters=128, kernel_size=3, activation='relu', padding='same'), # 卷积
        keras.layers.MaxPooling2D(pool_size=2),                                             # 池化

        # 为全连接层准备
        keras.layers.Flatten(), # 拉长数据，相当于将 立体三维特征图 变成 一维向量

        keras.layers.Dense(256, activation='relu'), # 得到 512个隐形特征，容易出现过拟合
        keras.layers.Dropout(0.5),  # （解决过拟合）感受野，每次忽略一半神经元 0.5 rate to drop
        # 猫狗分类二分类sigmoid就够了
        keras.layers.Dense(1, activation='sigmoid'),  # 相当于逻辑回归，映射到 0-1。多分类用activation='softmax'
    ])
    # model.summary() # 展示逻辑结构
    # 3.2 反向传播：配置训练器
    model.compile(loss='binary_crossentropy',
                  optimizer=Adam(lr=1e-4),
                  metrics=['acc'])  # 由于是二分类，使用准确率作为评估标准

    # 3.3 训练模型
    # fit 把所有数据读进内存；fit_generator 相当于生成器，动态生产所需 batch 数据
    history = model.fit(
        train_generator,   # 训练数据生成器
        steps_per_epoch=100,  # 每次迭代生成100个数据，至少需要跑 100次才能遍历完2000个数据。2000 images = batch_size * steps
        epochs=60,  # 迭代20次（每个 batch 是 20 个数据）
        validation_data=validation_generator,  # 目标数据生成器
        validation_steps=50,  # 1000 images = batch_size * steps
        verbose=2)

    # 4. 模型评估
    # 4.1 评估器评估

    # 4.2 画图展示
    acc = history.history['acc'] # 训练数据准确率
    val_acc = history.history['val_acc'] # 验证数据准确率
    loss = history.history['loss'] # 训练集 loss
    val_loss = history.history['val_loss'] # 验证集 loss

    epochs = range(len(acc))

    plt.plot(epochs, acc, 'bo', label='Training accuracy')  # 训练集精确率数据，点图例
    plt.plot(epochs, val_acc, 'b', label='Validation accuracy')  # 测试集精确率数据，线图例
    plt.title('Training and validation accuracy')  # 精确率图表标题

    plt.figure()
    plt.show()  # 展示图表

    plt.plot(epochs, loss, 'bo', label='Training Loss')  # 训练集损失值数据，点图例
    plt.plot(epochs, val_loss, 'b', label='Validation Loss')  # 测试集损失值数据，线图例
    plt.title('Training and validation loss')  # 损失值图表标题
    plt.legend()  # 添加图例

    plt.show()  # 展示图表



# # 5. 模型评估结果分类
# def predictModel(model, output_model_file):
#
#     # load the weights of model
#     model.load_weights(output_model_file)
#
#     os.makedirs('./save', exist_ok=True)
#     os.makedirs('./save/cat', exist_ok=True)
#     os.makedirs('./save/dog', exist_ok=True)
#
#     test_dir = './test/'  # 1-12500.jpg
#     for i in range(1, 12500):
#         img_name = test_dir + '{}.jpg'.format(i)
#         img = cv2.imread(img_name)
#         img = cv2.resize(img, (width, height))
#         img_arr = img / 255.0
#         img_arr = img_arr.reshape((1, width, height, 3))
#         pre = model.predict(img_arr)
#         if pre[0][0] > pre[0][1]:
#             cv2.imwrite('./save/cat/' + '{}.jpg'.format(i), img)
#             print(img_name, ' is classified as Cat.')
#         else:
#             cv2.imwrite('./save/dog/' + '{}.jpg'.format(i), img)
#             print(img_name, ' is classified as Dog.')



